import json
import os
from collections import OrderedDict
from glob import glob
from pathlib import Path

import cv2
from PIL import Image
from tqdm import tqdm

from ffpp.preprocessing import face_detector
from option import parse_args


def get_video_paths(original_path):
    paths = []
    paths.extend(glob(os.path.join(original_path, 'videos', "*.mp4")))
    return paths


def temp_func(x):
    return x


def get_video_frames(video):
    capture = cv2.VideoCapture(video)
    frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    frames = OrderedDict()
    for i in range(frames_num):
        capture.grab()
        success, frame = capture.retrieve()
        if not success:
            continue
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = Image.fromarray(frame)
        frame = frame.resize(size=[s // 2 for s in frame.size])
        frames[i] = frame
    return list(frames.keys()), list(frames.values())


def process_videos(video_paths, detector_cls: str):
    detector = face_detector.__dict__[detector_cls](landmarks=False)
    for video_path in tqdm(video_paths):
        video_id = os.path.splitext(os.path.basename(video_path))[0]
        parent_path = Path(video_path).parents[1]
        out_dir = os.path.join(parent_path, "boxes")
        output_json = os.path.join(out_dir, "{}.json".format(video_id))
        if os.path.exists(output_json):
            # print(output_json)
            continue
        indices, frames = get_video_frames(video_path)
        batches = [frames[i:i + detector.batch_size] for i in range(0, len(frames), detector.batch_size)]
        result = {}
        for j, frames in enumerate(batches):
            faces = detector.detect_faces(frames)
            for i, b in zip(indices, faces):
                result.update({int(j * detector.batch_size) + i: b})
        os.makedirs(out_dir, exist_ok=True)
        with open(os.path.join(out_dir, "{}.json".format(video_id)), "w") as f:
            json.dump(result, f)


def main():
    args = parse_args()
    if args.fake_type == '':
        path = os.path.join(args.root_dir, 'original_sequences', 'actors', args.compression_version)
    else:
        path = os.path.join(args.root_dir, 'manipulated_sequences', args.fake_type, args.compression_version)
    assert os.path.exists(path)
    video_paths = get_video_paths(path)
    print(len(video_paths))
    process_videos(video_paths, args.detector_type)


if __name__ == "__main__":
    main()
